#include "libtimers.h"


#define FAST_CHECK_BITS 4
#define FAST_CHECK_MASK ((1<<FAST_CHECK_BITS) - 1)

#define NULL_ID ((uint32_t)-1)
#define SLAB_SIZE_BITS 9
#define SLAB_SIZE (1<<SLAB_SIZE_BITS)
#define SLAB_SIZE_MASK (SLAB_SIZE-1)

typedef enum {
    LT_EFREE = 0,
    LT_EINUSE,
    LT_EINQUEUE,
} LT_STATE;

typedef struct {
    void *data;
    uint64_t expire;
    uint32_t prev;
    uint32_t next;
    uint32_t id;
    uint8_t state; 
} __attribute__((packed)) LT_TIMER;


typedef struct {
    struct {
        uint32_t head;
        uint32_t tail;
    } slots[256];

    uint64_t time_head;
    uint32_t granularity;
    uint8_t bitmaps[256 >> FAST_CHECK_BITS];
    uint16_t pos;
} LT_WHEEL;

#define DEFAULT_RET_SIZE 64

struct LT_TIMERPOOL{
    LT_WHEEL wheels[4];
    LT_TIMER **timers;
    uint32_t free_list;
    uint32_t slab_count;
    uint32_t slab_cap;
    uint32_t count;
    GET_TIME_FUNC time_func;
};


uint64_t get_now_time_ms()
{
    struct timeval tm;
    gettimeofday(&tm, NULL);
    return (uint64_t)tm.tv_sec * 1000 + tm.tv_usec / 1000;
}


#define get_by_id(pool, id) (&((pool)->timers[(id)>>SLAB_SIZE_BITS][(id) & SLAB_SIZE_MASK]))


/* normal linked-list operation, add an element to the tail  */
#define ADD_TO_LISTTAIL(node, anchor) {\
    (node)->prev = (anchor)->tail;\
    (node)->next = NULL_ID;\
    if ((anchor)->tail != NULL_ID) {\
        LT_TIMER *tm_tail = get_by_id(pool, (anchor)->tail);\
        tm_tail->next = (node)->id;\
    }\
    if ((anchor)->head == NULL_ID)\
        (anchor)->head = (node)->id;\
    (anchor)->tail = (node)->id;\
}

/* normal linked-list operation, delete an element */
#define DELETE_FROM_LIST(node, anchor) {\
    if ((anchor)->head == (node)->id)\
        (anchor)->head = (node)->next;\
    if ((anchor)->tail == (node)->id)\
        (anchor)->tail = (node)->prev;\
    if ((node)->prev != NULL_ID) {\
        LT_TIMER *tm_prev = get_by_id(pool, (node)->prev);\
        tm_prev->next = (node)->next;\
    }\
    if ((node)->next != NULL_ID) {\
        LT_TIMER *tm_next = get_by_id(pool, (node)->next);\
        tm_next->prev = (node)->prev;\
    }\
}


static LT_TIMER *new_timer(LT_TIMERPOOL *pool)
{
    LT_TIMER *timer;
    if (pool->free_list != NULL_ID) {
        timer = get_by_id(pool, pool->free_list);
        pool->free_list = timer->next;
        return timer;
    } 
    /* goes here means need to allocate a new slab */
    if (pool->slab_count == pool->slab_cap) {
        pool->slab_cap *= 2;    
        LT_TIMER **new_timers = (LT_TIMER**)malloc(sizeof(LT_TIMER*) * pool->slab_cap);
        for(size_t i = 0; i < pool->slab_count; i++) {
            new_timers[i] = pool->timers[i];
        }
        free(pool->timers);
        pool->timers = new_timers;
    }
    LT_TIMER *timer_slab = (LT_TIMER*)malloc(sizeof(LT_TIMER) * SLAB_SIZE);
    for(size_t i = 0; i < SLAB_SIZE; i++) {
        timer_slab[i].id = pool->slab_count * SLAB_SIZE + i;
        timer_slab[i].state = LT_EFREE;
        /* attach linked list */
        timer_slab[i].next = NULL_ID;
        if (i > 0)
            timer_slab[i - 1].next = timer_slab[i].id;
    }

    pool->timers[pool->slab_count++] = timer_slab;
    pool->free_list = timer_slab[1].id;
    return timer_slab; //&timer_slab[0]
}


static void free_timer(LT_TIMERPOOL *pool, LT_TIMER *timer)
{
    timer->next = pool->free_list;
    pool->free_list = timer->id;
}


/* run all timer events in a linked-list, which is attached on a slot position */
static void run_list(LT_TIMERPOOL *pool,
                    uint16_t pos, 
                    LT_TIMER_EVENT **events, 
                    int *events_count,
                    int *events_limit)
{
    uint32_t tid = pool->wheels[0].slots[pos].head;
    /* iterate the list */
    while(tid != NULL_ID) {
        LT_TIMER *timer = get_by_id(pool, tid);
        tid = timer->next;
        (*events)[*events_count].timer_id = timer->id;
        (*events)[*events_count].data = timer->data;
        timer->state = LT_EINUSE;
        (*events_count)++;
        pool->count--;
        if (*events_count == *events_limit) {
            /* enlarge events pool for return */
            *events_limit *= 2;
            LT_TIMER_EVENT *new_events = (LT_TIMER_EVENT*)malloc(sizeof(*new_events) * (*events_limit));
            for(int i = 0; i < *events_count; i++) 
                new_events[i] = (*events)[i];
            free(*events);
            *events = new_events;
        }
    }

    pool->wheels[0].slots[pos].head = NULL_ID;
    pool->wheels[0].slots[pos].tail = NULL_ID;
    return;
}


/* move the timers in a slot to the lower level
   and set them to the corresponding position by expire time */
static uint32_t dispatch_timers(LT_TIMERPOOL *pool, 
                                LT_WHEEL *wheel,
                                uint16_t pos, 
                                LT_WHEEL *wheel_child) 
{
    /* clear lower wheel's fast check bits */
    for(size_t i = 0; i < (256 >> FAST_CHECK_BITS); i++)
        wheel_child->bitmaps[i] = 0;

    wheel_child->pos = 0;
    /* base time of lower child is recalculated */
    wheel_child->time_head = wheel->time_head + pos * wheel->granularity;

    uint32_t tid = wheel->slots[pos].head;
    uint32_t count = 0;
    while(tid != NULL_ID) {
        LT_TIMER *timer = get_by_id(pool, tid);
        tid = timer->next;
        /* get its position by expire */
        uint16_t pos_child = (timer->expire - wheel_child->time_head)
                            / wheel_child->granularity;
        ADD_TO_LISTTAIL(timer, &wheel_child->slots[pos_child]);
        wheel_child->bitmaps[pos_child >> FAST_CHECK_BITS] = 1;
        count++;
    }
    /* clear the current slot */
    wheel->slots[pos].head = NULL_ID;
    wheel->slots[pos].tail = NULL_ID;

    /* return the number of moved timers */
    return count;
}


/* if a wheel's current position goes to the end, move timers from upper wheel here*/
static uint32_t shift_list(LT_TIMERPOOL *pool, int level, uint64_t now)
{
    LT_WHEEL *wheel = &pool->wheels[level];
    LT_WHEEL *wheel_child = &pool->wheels[level-1];

    /* walk the slots one by one until position got aligned */
    for(; wheel->pos & FAST_CHECK_MASK; wheel->pos++) {
        if (wheel->time_head + wheel->pos * wheel->granularity > now) 
            return dispatch_timers(pool, wheel, wheel->pos, wheel_child);

        /* try to move timers, though maybe no timer is moved */
        uint32_t shifted = dispatch_timers(pool, wheel, wheel->pos, wheel_child);
        wheel->pos++;
        if (shifted)
            return shifted;
    }

    for(;;) {
        if (wheel->time_head + wheel->pos > now)  
            return dispatch_timers(pool, wheel, wheel->pos, wheel_child);

        if (wheel->pos == 256) {
            if (level == 3) {
                /* top level is specially treated, since no timer moved to it */
                wheel->pos = 0;
                wheel->time_head += 256 * wheel->granularity;
            } else {
                if (!shift_list(pool, level + 1, now)) {
                    /* got no moved timers, try to shift_list immediately */
                    //wheel->pos = 256;
                    //continue;
                }
            }
        }
 
        if (wheel->bitmaps[(wheel->pos) >> FAST_CHECK_BITS] || level == 3) {
            /* actually fast check at level 3 is useless 
               because bitmap in level 3 was never set */
            do {
                if (wheel->time_head + wheel->pos > now) {
                    /* return if surpassed now */
                    wheel_child->pos = 0;
                    wheel_child->time_head = wheel->time_head
                        + wheel->pos * wheel->granularity;
                    return 0;
                }

                uint32_t shifted = dispatch_timers(pool, wheel, wheel->pos, wheel_child);
                wheel->pos++;
                if (shifted)
                    return shifted;
            } while(wheel->pos & FAST_CHECK_MASK);
        } else
            /* skip some slots if no timers in these slots */
            wheel->pos += (1 << FAST_CHECK_BITS);
    }
}


static int del_timer(LT_TIMERPOOL *pool, LT_TIMER *timer)
{
    for(size_t i = 0; i < 4; i++) {
        /* from bottom wheel to top wheel, check which one it should belongs to */
        LT_WHEEL *wheel = &pool->wheels[i];
        if (timer->expire < wheel->time_head + wheel->granularity * 256) {
            uint16_t pos = (timer->expire - wheel->time_head) / wheel->granularity;
            DELETE_FROM_LIST(timer, &wheel->slots[pos]);
            pool->count--;
            return 0;
        }
    }
    return -1;
}

static int set_timer(LT_TIMERPOOL *pool, LT_TIMER *timer, void *data, uint32_t interval)
{
    uint64_t now = pool->time_func();
    timer->data = data;
    timer->expire = now + interval;
    for(size_t i = 0; i < 4; i++) {
        /* from bottom wheel to top wheel, check which one it should belongs to */
        LT_WHEEL *wheel = &pool->wheels[i];
        if (timer->expire < wheel->time_head + wheel->granularity * 256) {
            /* got! set into the proper position */
            uint16_t pos = (timer->expire - wheel->time_head) / wheel->granularity;
            ADD_TO_LISTTAIL(timer, &wheel->slots[pos]);
            timer->state = LT_EINQUEUE;
            wheel->bitmaps[pos >> FAST_CHECK_BITS] = 1;
            pool->count++;
            return timer->id;
        }
    }
    /* no wheel found! */
    assert(0);
    return NULL_ID;
}


LT_TIMERPOOL *lt_new(GET_TIME_FUNC func)
{
    LT_TIMERPOOL *pool = (LT_TIMERPOOL*)malloc(sizeof(*pool));
    uint64_t now = 0;

    /* time function, use default unless speicifed */
    if ((pool->time_func = func) == NULL)
        pool->time_func = get_now_time_ms;

    now = pool->time_func();

    /* wheels have granularities of 1, 256, 65536, 16777216 */
    /* wheels[0] now points to pos 0, others points to pos 1*/
    pool->wheels[0].granularity = 1;
    pool->wheels[0].pos = 0;

    for(int i = 0; i < 4; i++) {
        LT_WHEEL *wheel = &pool->wheels[i];
        if (i > 0) {
            wheel->granularity = pool->wheels[i - 1].granularity * 256;
            wheel->pos = 1;
        }
        
        /* base time of current wheel is now */
        wheel->time_head = now;

        /* all wheels has no timer linked on it */
        for(size_t j = 0; j < 256; j++) 
            wheel->slots[j].head = wheel->slots[j].tail = NULL_ID;
    }

    /* timer objects are stored in 2-D arrays,
       with fixed size of 'SLAB_SIZE' objects on 2nd dimension */
    pool->slab_count = 0;
    pool->count = 0;
    pool->slab_cap = 8;
    pool->timers = (LT_TIMER**)malloc(sizeof(LT_TIMER*) * pool->slab_cap);

    /* objects are stored in arrays, also double-linked with each other */
    pool->free_list = NULL_ID;

    return pool;
}

void lt_destroy(LT_TIMERPOOL *pool) 
{
    /* free all arrays */
    for(uint32_t i = 0; i < pool->slab_count; i++) 
        free(pool->timers[i]);

    free(pool->timers);
    free(pool);
}

uint32_t lt_timer_open(LT_TIMERPOOL *pool)
{
    LT_TIMER *timer = new_timer(pool);
    timer->state = LT_EINUSE;
    return timer->id;
}

int lt_timer_close(LT_TIMERPOOL *pool, uint32_t timer_id)
{
    if (timer_id >= pool->slab_count * SLAB_SIZE)
        return -1;

    LT_TIMER *timer = get_by_id(pool, timer_id);
    if (timer->state == LT_EFREE) {
        return -1;
    } else if (timer->state == LT_EINUSE) {
        timer->state = LT_EFREE;
        free_timer(pool, timer);
        return 0;
    } else if (timer->state == LT_EINQUEUE) {
        del_timer(pool, timer);
        timer->state = LT_EFREE;
        free_timer(pool, timer);
        return 0;
    } else {
        return -1;
    }
}


int lt_set_timer(LT_TIMERPOOL *pool, uint32_t timer_id, void *data, uint32_t interval)
{
    if (timer_id >= pool->slab_count * SLAB_SIZE)
        return -1;

    LT_TIMER *timer = get_by_id(pool, timer_id);
    if (timer->state == LT_EFREE) {
        return -1;
    } else if (timer->state == LT_EINQUEUE) {
        del_timer(pool, timer);
    }

    return set_timer(pool, timer, data, interval);
}

int lt_cancel_timer(LT_TIMERPOOL *pool, uint32_t timer_id)
{
    if (timer_id >= pool->slab_count * SLAB_SIZE)
        return -1;

    LT_TIMER *timer = get_by_id(pool, timer_id);
    return del_timer(pool, timer);
}


/* like what linux epoll did, return the time up events */
int lt_poll(LT_TIMERPOOL *pool, LT_TIMER_EVENT **events)
{
    uint64_t now = pool->time_func();

    LT_WHEEL *wheel = &pool->wheels[0];

    int events_count = 0;
    int events_limit = DEFAULT_RET_SIZE;
    *events = (LT_TIMER_EVENT *)malloc(sizeof(**events) * events_limit);

    /* walk the slots one by one until position got aligned */
    for(; wheel->pos & FAST_CHECK_MASK; wheel->pos++) {
        if (wheel->time_head + wheel->pos > now)
            /* return if surpassed now */
            return events_count;

            /* do timer events on this slot */
            run_list(pool, wheel->pos, events, &events_count, &events_limit);
    }

    for(;;) {
        if (wheel->time_head + wheel->pos > now) 
            return events_count;

        if (wheel->pos == 256) {
            /* walked to the end of wheel */
            if (!shift_list(pool, 1, now)) {
                /* got no moved timers, try to shift_list immediately */
                //wheel->pos = 256;
                //continue;
            }
        }
 
        if (wheel->bitmaps[wheel->pos >> FAST_CHECK_BITS]) {
            do {
                if (wheel->time_head + wheel->pos > now)
                    return events_count;

                /* do timer events on this slot */
                run_list(pool, wheel->pos, events, &events_count, &events_limit);
                wheel->pos++;
            } while(wheel->pos & FAST_CHECK_MASK);
        } else 
            wheel->pos += (1<<FAST_CHECK_BITS);
    }

    return events_count;
}

void lt_free_events(LT_TIMERPOOL *pool, LT_TIMER_EVENT *events)
{
    free(events);
}

int lt_timer_size(LT_TIMERPOOL *pool)
{
    return pool->count;
}


